import httpx
import m3u8
import os
import re


class M3U8(object):
    @classmethod
    def download_video(cls, m3u8_url, save_dir, filename):
        # m3u8_url = 'https://video.fedbf.com/20230719/M2Y1NWRlM2/123818/720/hls/encrypt/index.m3u8'
        # m3u8_url = 'https://video2.tangrenfuli.com/video/20210921/c2d8017defe31bfb7e411eaf4c6a951d/index.m3u8'
        # output_file = 'output.mp4'

        # 解析m3u8文件
        playlist = m3u8.load(m3u8_url)
        # 获取视频流列表
        video_streams = playlist.segments
        if len(video_streams) <= 0:
            return False

        # 得到后缀，并添加到路径后面
        ts_url = video_streams[0].uri  # 获取每个片段的URL
        suffix = "mp4"
        ts_url_items = ts_url.split(".")
        if len(ts_url_items) > 1:
            suffix = ts_url_items[-1]
        save_path = f"{save_dir}/{filename}.{suffix}"

        # 创建输出目录（如果不存在）
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        count = len(video_streams)
        print(f"下载:{m3u8_url}, 分片数:{count}, 保存到: {save_path} ....")
        # 下载视频流并合并为单个文件
        with open(save_path, 'wb') as output_file:
            for i in range(len(video_streams)):
                segment = video_streams[i]
                ts_url = segment.uri  # 获取每个片段的URL
                http_index = ts_url.find("http")
                if http_index != 0:
                    match = re.search(r'/([^/]+)$', m3u8_url)
                    if match:
                        # file_part = match.group(1)
                        ts_url = m3u8_url[:match.start()] + "/" + ts_url + m3u8_url[match.end():]
                print(f"下载分片({i+1}/{count}): {ts_url}...")
                ts_response = httpx.get(ts_url)  # 发送GET请求获取片段数据

                if ts_response:
                    ts_data = ts_response.content  # 获取片段数据内容
                    if ts_data:
                        output_file.write(ts_data)  # 将片段数据写入输出文件

        print(f"下载:{m3u8_url}, 保存到: {save_path} 完成")


if __name__ == '__main__':
    M3U8.download_video(
        m3u8_url="https://video2.tangrenfuli.com/video/20210921/c2d8017defe31bfb7e411eaf4c6a951d/index.m3u8",
        save_dir="/Users/yuanzhihua/study/source/spider",
        filename="ccccc"
    )
